# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon

import logging

from absl import app, flags
from importlib import import_module

from coin.strategy.interval.interval_dumper import universe_btc, get_dump_root_dir
from coin.strategy.interval.util import get_df_from_total, load_recent_df
from coin.strategy.interval.plot_alpha import print_matplotlib, plotting_init, plotting
from coin.strategy.interval.alphas.taekwon_fit import cross_fit
from coin.strategy.interval.transforms import identity, logistic_normalize
from coin.strategy.interval.model import Transform, LinearModel

from flask import Flask, render_template
from flask import send_file, request

from gevent.server import _tcp_listener
from gevent import pywsgi

import coin.strategy.mm.tool.archive_base as abase


def plot_single_alpha(products, alpha_trans, y_feature, ith, df):
  plotting_init()
  name = '_'.join(products)
  x = alpha_trans.eval(df)
  t = df['global_timestamp']
  y = df[f'Y/{products[0]}/{y_feature}']
  p = df[f'{products[0]}/b.mid_p']
  return print_matplotlib(x, y, t, p, name, desc=alpha_trans.alpha.name)


flaskapp = None


def plot_single_image(products, alpha_trans_lst, y_feature, ith):
  img = plot_single_alpha(products, alpha_trans_lst[int(ith)], y_feature, ith, TOTAL_DF)
  return send_file(img, mimetype='image/png')


def load_module(module_path, products):
  p, m = module_path.rsplit('.', 1)
  alpha_trans_lst = getattr(import_module(p), m)(products)
  return alpha_trans_lst


TOTAL_DF = load_recent_df()
# warning: hack!
SYMBOLS = [col.split("/")[0] for col in TOTAL_DF.columns if col.endswith("/b.ask0_p")]

flaskapp = Flask(__name__)


@flaskapp.route('/plot/<author>/<profile>/<ith>', methods=['GET'])
def plot_specific_image(author, profile, ith):
  args = request.args
  products_str = args.get('products', 'bitflyer-fxbtcjpy')
  y_feature = args.get('y', 'vwap_ret.1m')
  print(products_str)
  products = products_str.split(',')
  alpha_trans_lst = load_module(f"coin.strategy.interval.profile.{author}.{profile}", products)
  return plot_single_image(products, alpha_trans_lst, y_feature, ith)


@flaskapp.route('/plot/<author>/<profile>', methods=['GET'])
def plot(author, profile):
  args = request.args
  products_str = args.get('products', 'bitflyer-fxbtcjpy')
  products = products_str.split(',')
  y_feature = args.get('y', 'vwap_ret.1m')

  alpha_trans_lst = load_module(f"coin.strategy.interval.profile.{author}.{profile}", products)
  num_images = len(alpha_trans_lst)
  print(alpha_trans_lst)
  return render_template('alpha.html',
                         request=request,
                         symbols=SYMBOLS,
                         products=products,
                         num_images=num_images)


@flaskapp.route('/ridge_fit/<author>/<profile>', methods=['GET'])
def ridge_fit(author, profile):
  args = request.args
  products_str = args.get('products', 'bitflyer-fxbtcjpy')
  filename = args.get('filename', f'ridge_fit_{author}_{profile}')
  path = f'coin/strategy/interval/model/{filename}.pkl'
  products = products_str.split(',')
  y_feature = args.get('y', 'vwap_ret.1m')
  y = TOTAL_DF[f'Y/{products[0]}/{y_feature}'].fillna(0)

  alpha_trans_lst = load_module(f"coin.strategy.interval.profile.{author}.{profile}", products)
  model = LinearModel(alpha_trans_lst)
  _, fit_scores = model.ridge_fit(TOTAL_DF, y)
  model.save(path)
  template_str = "<br/><br/>IS rsq:{}, OOS req: {}</br>".format(fit_scores[0], fit_scores[1])
  for coef, alpha_trans in model.get_coef_alpha_trans():
    template_str += "<br/>Coef:{}, Alpha: {}".format(coef, alpha_trans.alpha.name)
  return f"Success saving on {path}!" + template_str


@flaskapp.route('/ridge_plot/<author>/<profile>')
def ridge_plot(author, profile):
  args = request.args
  products_str = args.get('products', 'bitflyer-fxbtcjpy')
  products = products_str.split(',')
  y_feature = args.get('y', 'vwap_ret.1m')

  alpha_trans_lst = load_module(f"coin.strategy.interval.profile.{author}.{profile}", products)
  model = LinearModel(alpha_trans_lst)
  y = TOTAL_DF[f'Y/{products[0]}/{y_feature}'].fillna(0)
  t = TOTAL_DF['global_timestamp']
  p = TOTAL_DF[f'{products[0]}/b.mid_p']
  y_pred, fit_scores = model.ridge_fit(TOTAL_DF, y)
  desc = f"{str(products)}\nIS_score:{fit_scores[0]}, OOS_score:{fit_scores[1]}"

  img_path = plotting(y_pred, y, t, p, f'Ridge_Plot_{author}_{profile}', desc=desc)
  return render_template('index.html', symbols=[], plot_img_paths=[img_path])


@flaskapp.route('/<author>/<profile>', methods=['GET'])
def plot_specific(author, profile):
  base = request.args.get("base", "okex-btcusd-quarter")
  ref = request.args.get("ref", "bitmex-xbtusd")
  return load_module(f"coin.strategy.interval.profile.{author}.{profile}", request, base, ref)


def main(argv):
  if True:  # set False to run in debug mode
    listener = _tcp_listener(('0.0.0.0', 5883))

    def serve_forever():
      server = pywsgi.WSGIServer(listener, flaskapp)
      server.serve_forever()

    from multiprocessing import Process
    number_of_processes = 5
    for i in range(number_of_processes):
      print(f"server {i}")
      Process(target=serve_forever).start()
  else:
    flaskapp.run(debug=True, port=5882, host='0.0.0.0')


def main_headless(argv):
  load_module("coin.strategy.interval.profile.taekwon.greats", {})


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()

  flags.FLAGS.exchange = ""
  flags.FLAGS.market_type = ""

  logging.basicConfig(level='DEBUG')
  app.run(main)
  # app.run(main_headless)
